package com.runmyai.aitools.text;

import com.runmyai.aitools.utils.Constants;
import org.springframework.util.StringUtils;

import java.util.*;
import java.util.regex.Pattern;
import java.util.stream.Collectors;
import java.util.stream.Stream;

/**
 * 处理各类型带序号/不带序号的列表数据。
 * 建议在提示词中加上一条：
 *
 *   格式要求：
 *   1. 带上序号
 *
 * 能获得更好的稳定性
 *
 * @author wry
 */
public class SequenceTextParser {

    /**
     * markdown中可用来表示多行数据的前缀。
     * 有序列表： 1.
     * 无序列表： -
     * 无序列表： *
     * 分类+换行： xxx：\n
     */
    private static final Pattern SEQUENCE_PATTERN = Pattern.compile(Constants.SEQUENCE_REGEX);
    private static final Pattern UNORDERED_PATTERN = Pattern.compile(Constants.UNORDERED_REGEX);
    private static final Pattern STAR_SEQUENCE_PATTERN = Pattern.compile(Constants.STAR_SEQUENCE_REGEX);
    private static final Pattern COLON_SEQUENCE_PATTERN = Pattern.compile(Constants.COLON_SEQUENCE_REGEX);

    public static Optional<List<String>> parse(String text) {
        List<String> combined = new ArrayList<>();

        Stream<String> stream = Arrays.stream(text.split("\\n\\n"))
                .filter(StringUtils::hasText)
                .map(String::trim);

        Map<Boolean, List<String>> result = stream.collect(Collectors.partitioningBy(str -> SEQUENCE_PATTERN.matcher(str).find()));

        // 获取初次匹配通过的字符串列表
        List<String> filteredList = result.getOrDefault(true, Collections.emptyList());

        // 对未通过的字符串列表继续进行匹配,得到第二次结果
        Map<Boolean, List<String>> secondResult = result.getOrDefault(false, Collections.emptyList()).stream()
                .collect(Collectors.partitioningBy(str -> UNORDERED_PATTERN.matcher(str).find()));
        filteredList.addAll(secondResult.getOrDefault(true, Collections.emptyList()));

        // 对第二次未通过的字符串列表继续进行匹配,得到第三次结果
        Map<Boolean, List<String>> thirdResult = secondResult.getOrDefault(false, Collections.emptyList()).stream()
                .collect(Collectors.partitioningBy(str -> STAR_SEQUENCE_PATTERN.matcher(str).find()));
        filteredList.addAll(thirdResult.getOrDefault(true, Collections.emptyList()));

        // 对最终未通过的字符串列表进行处理. 如果有  xxx：\n  则截取剩余的字符串
        Map<Boolean, List<String>> fourthResult = thirdResult.getOrDefault(false, Collections.emptyList()).stream()
                .collect(Collectors.partitioningBy(str -> COLON_SEQUENCE_PATTERN.matcher(str).find()));

        // 处理 xxxA： 开头的文本
        // 无法处理 xxxA：xx \n xxxB：yy \n xxxC：
        fourthResult.getOrDefault(true, Collections.emptyList()).stream()
                .map(str -> str.replaceFirst(Constants.COLON_SEQUENCE_REGEX, ""))
                .forEach(filteredList::add);

        //处理前后都是中文引号的内容
        fourthResult.getOrDefault(false, Collections.emptyList()).stream()
                .filter(str -> str.startsWith("“") && str.endsWith("”"))
                .forEach(filteredList::add);

        // 将最终的单行文本用TextPostProcessor处理
        filteredList.stream().flatMap(input -> Arrays.stream(input.split("\\n"))
                        .map(TextPostProcessor::clean))
                .filter(StringUtils::hasText)
                .forEach(combined::add);

        return Optional.ofNullable(combined);
    }
}
